#include "ArmorDetector.h"

ArmorDetector::ArmorDetector(const string &filename)
{
    FileStorage FS(filename, FileStorage::READ);
    assert(FS.isOpened());
    FS["camera_matrix"] >> this->cameraMatrix;
    FS["distortion_coefficients"] >> this->distCoeff;
    FS["track_frames"] >> this->TRACK_FRAMES;
    this->save_armors = vector<vector<Armor_ptr>>(this->TRACK_FRAMES);
}

void ArmorDetector::detectArmor(Mat &src_img)
{
    this->current_armors.clear();
    src_img.copyTo(this->draw_img);
    single_blob.clear();
    vector<vector<Point>> contours;
    vector<light_blobs> tmp_blob;
    if (this->backgound.empty())
    {
        this->backgound = Mat::zeros(src_img.size(), CV_8UC3);
    }
    if (src_img.empty())
    {
        cout << "no frames!!" << endl;
        return;
    }

    img_proc(src_img, contours);

    if (contours.size() == 0)
    {
        return;
    }

    find_lightblob(contours, single_blob);
    match_lightblob(single_blob, tmp_blob);

    if (tmp_blob.size() != 0)
    {
        for (auto i : tmp_blob)
        {
            Armor_ptr armor_o(new Armor);
            armor_o->blobs = i;
            armor_o->get_center();
            armor_o->record_time = getTickCount();
            armor_o->pnp_data = armor_o->get_pnpdata(armor_o->type,armor_o->blobs,this->cameraMatrix,this->distCoeff);
            current_armors.push_back(armor_o);
            ArmorDetector::draw_blob(i.left->corners, i.right->corners, draw_img);
            ArmorDetector::draw_blob(i.left->corners, i.right->corners, backgound);
            ArmorDetector::draw_blob(i.left->corners, i.right->corners, src_img);
        }
    }

    if (save_num <= this->TRACK_FRAMES) //保存TRACK_FRAMES帧//
    {
        for (auto i : current_armors)
        {
            this->save_armors[this->TRACK_FRAMES - this->save_num].push_back(i);
        }
        ++this->save_num;
    }
    else
    {
        for (int i = 0; i < current_armors.size(); ++i)
        {
            for (int j = 0; j < save_armors.size(); ++j)
            {
                current_armors[i]->Search_last(save_armors[j], j);
                if (current_armors[i]->isMatched)
                {
                    if (!current_armors[i]->hasId)
                    {
                        current_armors[i]->id = id;
                        current_armors[i]->hasId = true;
                        ++id;
                        break;
                    }
                }
            }
        }
    }

    filter_repeated(current_armors);
    if (save_num > this->TRACK_FRAMES)
    {
        for (int i = this->TRACK_FRAMES - 1; i >= 1; --i) //数组后推//
        {
            save_armors[i] = save_armors[i - 1];
        }
        save_armors[0] = current_armors;
    }
    stringstream s;
    cout << current_armors.size() << endl;
    for (auto i : current_armors)
    {
        if (i->hasId)
        {
            // cout << "id:" << i->id << endl;
            string num;
            s << i->id;
            s >> num;
            circle(draw_img, i->armor_center, 2, Scalar(255, 255, 255));
            putText(draw_img, num, i->blobs.left->corners[0], FONT_HERSHEY_COMPLEX, 0.3, cv::Scalar(255, 255, 255));
        }
    }

    if (current_armors.size() == 0)
    {
        return;
    }
    // for (auto i : current_armors)
    // {
    //     i->pnp_data = i->get_pnpdata(i->type, i->blobs, this->cameraMatrix, this->distCoeff);
    // }
    sort(current_armors.begin(), current_armors.end(), [](Armor_ptr armor1, Armor_ptr armor2) { return armor1->pnp_data.distance < armor1->pnp_data.distance; });
    // cout << "distance:" << current_armors[0]->pnp_data.distance << endl;
    // cout << "yaw" << current_armors[0]->relative_yaw << endl;
    // cout << "picth" << current_armors[0]->relative_picth << endl;
    current_armors[0]->get_center_distance();
    if (current_armors[0]->kalmanInited && current_armors[0]->kalman_center.x < 640 && current_armors[0]->kalman_center.y < 480)
    {
        circle(draw_img, current_armors[0]->kalman_center, 3, Scalar(0, 0, 255));     
    }
}

void ArmorDetector::img_proc(Mat &src_img, vector<vector<Point>> &contours)
{
    // Mat gray_img, bin_img;

    // if (src_img.channels() == 3)
    // {
    //     // cvtColor(src_img, gray_img, CV_BGR2GRAY);
    //     std::vector<cv::Mat> channals_mat;
    //     cv::split(src_img, channals_mat);
    //     gray_img = channals_mat[0];
    // }
    // else
    //     gray_img = src_img;

    // threshold(gray_img, src_img,
    //           //COLOR_MODE: 0-红 1-蓝
    //           40,255, THRESH_BINARY);
    cvtColor(src_img, src_img, COLOR_BGR2HSV);
    inRange(src_img, Scalar(hmin, smin, vmin), Scalar(hmax, smax, vmax), src_img);
    findContours(src_img, contours, RETR_EXTERNAL, CHAIN_APPROX_NONE);
}

void ArmorDetector::find_lightblob(vector<vector<Point>> &contours, vector<Blob_ptr> &single_blob)
{
    vector<Point> blob_corners1(2);
    vector<Point> blob_corners2(2);
    RotatedRect ellipse0;
    RotatedRect ellipse1;
    int x;
    int y;
    int macth_num;
    int min_num = -1;
    int width; //装甲片的长和宽//
    int height;
    int x_diff; //两个轮廓的水平距离//
    double scale3;
    float current_scale;
    float min_scale = 1.9;
    float rect_scale0;
    float rect_scale1;
    float rect_scale3;

    vector<Rect> boundRect(contours.size()); //矩形逼近每个轮廓//
    for (int i = 0; i < contours.size(); i++)
    {
        boundRect[i] = boundingRect(Mat(contours[i]));
    }

    for (int i = 0; i < contours.size(); i++) //遍历每个轮廓
    {
        rect_scale0 = (float)boundRect[i].height / boundRect[i].width;
        if (rect_scale0 <= 1 || contours[i].size() < 15)
        {
            continue;
        }

        Blob_ptr tool_blob(new LightBlob);
        tool_blob->blob_contour = contours[i];
        single_blob.push_back(tool_blob);
    }
    for (auto i : single_blob)
    {
        find_corners(i);
        i->color_distance0 = this->draw_img.at<Vec3b>(i->corners[0])[0] - this->draw_img.at<Vec3b>(i->corners[0])[2];
        i->color_distance1 = this->draw_img.at<Vec3b>(i->corners[1])[0] - this->draw_img.at<Vec3b>(i->corners[1])[2];
        i->angle = atan2(i->corners[0].y - i->corners[1].y, i->corners[0].x - i->corners[1].x) / 3.1415 * 180;
        i->blob_length = distance(i->corners);
    }
    single_blob.erase(remove_if(single_blob.begin(), single_blob.end(), [](Blob_ptr i) { return i->color_distance0 <= 30 || i->color_distance1 <= 30; }), single_blob.end());
}

void ArmorDetector::match_lightblob(vector<Blob_ptr> &single_blob, vector<light_blobs> &matched_blobs)
{

    double min_scale = 8;
    int min_index = -1;
    for (int i = 0; i < single_blob.size(); i++)
    {
        vector<Blob_ptr> candidate_blob;
        if (single_blob[i]->isMatched)
        {
            continue;
        }
        min_scale = 8;
        min_index = -1;
        double last_y_distance;
        for (int j = 0; j < single_blob.size(); j++)
        {
            if (single_blob[j]->isMatched)
            {
                continue;
            }
            if (i == j || single_blob[i]->blob_contour.size() < 15 || single_blob[j]->blob_contour.size() < 15)
            {
                continue;
            }

            double length_i = single_blob[i]->blob_length;
            double length_j = single_blob[j]->blob_length;
            double length_scale = length_i / length_j >= 1 ? length_i / length_j : length_j / length_i;
            if (fabs(single_blob[i]->angle - single_blob[j]->angle) >= 9 || length_scale >= 1.4)
            {
                continue;
            }

            double Armor_width = length_i;
            double Armor_length = distance(single_blob[i]->corners[0], single_blob[j]->corners[0]);
            double Armor_scale = Armor_length / Armor_width;
            if (!(Armor_scale <= 5.5 && Armor_scale >= 1))
            {
                continue;
            }

            double up_distance = fabs(single_blob[i]->corners[0].y - single_blob[j]->corners[0].y) + 1; //加1防止其等于零//
            double down_distance = fabs(single_blob[i]->corners[1].y - single_blob[j]->corners[1].y) + 1;
            if (down_distance >= (length_i > length_j ? length_j : length_i))
            {
                continue;
            }

            double up_dowm_scale = (up_distance / down_distance) >= 1 ? (up_distance / down_distance) : (down_distance / up_distance); //角度差接近时用上下距离差的比决定//
            double angle_scale = fabs(single_blob[i]->angle - single_blob[j]->angle);

            double area_i = contourArea(single_blob[i]->blob_contour);
            double area_j = contourArea(single_blob[j]->blob_contour);

            if (area_i <= 20 || area_j <= 20)
            {
                continue;
            }

            if (angle_scale >= min_scale)
            {
                continue;
            }

            single_blob[j]->sum_scale_error = length_scale + up_dowm_scale + angle_scale - 3;
            candidate_blob.push_back(single_blob[j]);
            min_index = 0;
        }
        if (min_index != -1)
        {
            sort(candidate_blob.begin(), candidate_blob.end(), [](Blob_ptr blob1, Blob_ptr blob2) { return blob1->sum_scale_error < blob2->sum_scale_error; });
            light_blobs tool_light;
            tool_light.left = single_blob[i]->corners[0].x > candidate_blob[0]->corners[0].x ? candidate_blob[0] : single_blob[i];
            tool_light.right = single_blob[i]->corners[0].x <= candidate_blob[0]->corners[0].x ? candidate_blob[0] : single_blob[i];
            matched_blobs.push_back(tool_light);
            candidate_blob[0]->isMatched = true;
            single_blob[i]->isMatched = true;
        }
    }
}

void ArmorDetector::find_corners(Blob_ptr &single_blob)
{
    int uppest_y = 480;
    int lowest_y = 0;
    for (auto i : single_blob->blob_contour)
    {
        if (i.y < uppest_y)
        {
            uppest_y = i.y;
            single_blob->corners[0].x = i.x;
            single_blob->corners[0].y = i.y;
        }
    }
    for (auto i : single_blob->blob_contour)
    {
        if (i.y > lowest_y)
        {
            lowest_y = i.y;
            single_blob->corners[1].y = i.y;
            single_blob->corners[1].x = i.x;
        }
    }
}

void ArmorDetector::draw_blob(vector<Point2f> &blob_corners1, vector<Point2f> &blob_corners2, Mat &src_img)
{
    circle(src_img, blob_corners1[0], 2, Scalar(0, 255, 0));
    circle(src_img, blob_corners1[1], 2, Scalar(0, 255, 0));
    circle(src_img, blob_corners2[0], 2, Scalar(0, 255, 0));
    circle(src_img, blob_corners2[1], 2, Scalar(0, 255, 0));
    line(src_img, blob_corners1[0], blob_corners2[1], Scalar(255, 255, 255));
    line(src_img, blob_corners1[1], blob_corners2[0], Scalar(255, 255, 255));
}

double ArmorDetector::distance(vector<Point2f> &blob)
{
    return sqrt(pow(blob[0].x - blob[1].x, 2) + pow(blob[0].y - blob[1].y, 2));
}

double ArmorDetector::distance(Point2f a, Point2f b)
{
    return sqrt(pow(a.x - b.x, 2) + pow(a.y - b.y, 2));
}

void ArmorDetector::filter_repeated(vector<Armor_ptr> &current_armors)
{
    for (int i = 0; i < current_armors.size(); i++) //修正前几帧中有一个装甲板匹配了目前帧中多个装甲板的错误//
    {
        if (!current_armors[i]->hasId)
        {
            continue;
        }
        for (int j = 0; j < current_armors.size(); j++)
        {
            if (!current_armors[j]->hasId)
            {
                continue;
            }
            if (i == j)
            {
                continue;
            }
            if (current_armors[i]->id == current_armors[j]->id)
            {
                if (current_armors[i]->numberOf_frame < current_armors[j]->numberOf_frame)
                {
                    current_armors[j]->hasId = false;
                    for (int k = 0; k < save_armors.size(); k++)
                        current_armors[j]->Search_last(save_armors[k], k);
                }
                else
                {
                    current_armors[i]->hasId = false;
                    for (int k = 0; k < save_armors.size(); k++)
                        current_armors[i]->Search_last(save_armors[k], k);
                }
            }
        }
    }
}